import pickle
import numpy as np
import torch
import os

from DataSet.Dataset_DPN107 import Dataset_DPN107
from DataSet.Tools_DPN107 import load_DPN107_from_file,extract_DPN107_feature
from DataSet.Utils import mulity_scale_pairs

class Dataset_DPN107_Pair(Dataset_DPN107):

    def _random_sample_traing_item(self,batch=2048, random=512, hard_neg=512, maxnump=32, PYRAMID=[12,12], thr_ious=[0.7,0.3]):
        '''
        无穷无尽的取训练用的batch
        :param batch:
        :param random:
        :param hard_neg:
        :param maxnump:
        :param PYRAMID:
        :param thr_ious:
        :return:
        '''
        target_pos = batch-hard_neg-random
        pos_count = 0
        neg_count = 0
        rand_count = 0
        count = 0

        x = np.zeros((batch, 1, np.sum(PYRAMID), 2688))
        y = np.ones(batch)*-1

        for pos,vid in enumerate(self.vids):

            vitem = self.groundtruth[vid]

            if os.path.exists('/mnt/md1/Experiments/ExA_Test17/Features/dpn107_{}.pkl'.format(vid)) == False \
                    or vid not in self.proposals or len(self.proposals[vid]) == 0:
                continue

            feat = load_DPN107_from_file(vid)
            duration = vitem['duration']
            proposal = np.asarray(self.proposals[vid])

            max_ious = proposal[:,:,2].max(axis=1)
            max_ious_lv = proposal[:,:,2].argmax(axis=1)

            pos_indexes = np.nonzero(max_ious>thr_ious[0])[0]
            neg_indexes = np.nonzero(max_ious<thr_ious[1])[0]

            pos_pairs = []
            for i in range(len(pos_indexes)):
                ind = pos_indexes[i]
                lv = max_ious_lv[ind]
                pos_pairs += mulity_scale_pairs(proposal[ind,], [lv])

            np.random.shuffle(pos_pairs)
            for pair in pos_pairs[:min(target_pos,maxnump)]:
                if pos_count < target_pos and count < batch:
                    data = extract_DPN107_feature(feat, pair, duration, PYRAMID=PYRAMID)
                    if data is not None and data.shape[0]==np.sum(PYRAMID):
                        x[count, ...].flat = data.flat
                        y[count, ...] = 1
                        pos_count += 1
                        count += 1

            # sample random negative around around positive
            rand_pairs = []
            for i in range(len(pos_indexes)):
                ind = pos_indexes[i]
                lv = max_ious_lv[ind]
                rand_pairs += mulity_scale_pairs(proposal[ind,], [i for i in range(6) if i!=lv])
            np.random.shuffle(rand_pairs)

            for pair in rand_pairs[:min(random,maxnump)]:
                if rand_count<random and count < batch:
                    data = extract_DPN107_feature(feat, pair, duration, PYRAMID=PYRAMID)
                    if data is not None and data.shape[0] == np.sum(PYRAMID):
                        x[count, ...].flat = data.flat
                        y[count, ...] = 0
                        rand_count += 1
                        count += 1

            # sample hard negatives
            neg_pairs = []
            for i in range(len(neg_indexes)):
                ind = neg_indexes[i]
                neg_pairs += mulity_scale_pairs(proposal[ind,])
            np.random.shuffle(neg_pairs)

            for pair in neg_pairs[:min(hard_neg,maxnump)]:
                if neg_count < hard_neg and count < batch:
                    data = extract_DPN107_feature(feat, pair, duration, PYRAMID=PYRAMID)
                    if data is not None and data.shape[0]==np.sum(PYRAMID):
                        x[count, ...].flat = data.flat
                        y[count, ...] = 0
                        neg_count += 1
                        count += 1

            if count == batch:
                # 位移动一下
                self.vids = self.vids[pos:] + self.vids[:pos]
                break

        return x,y

    def _enum_every_test_item(self):
        # 一次出来一组6个和对应的VID
        for vid,proposals in self.proposals.items():

            vitem = self.groundtruth[vid]
            if os.path.exists('/mnt/md1/Experiments/ExA_Test17/Features/dpn107_{}.pkl'.format(vid)) == False or \
                    vid not in self.proposals or len(self.proposals[vid]) == 0:
                continue

            feat = load_DPN107_from_file(vid)
            duration = vitem['duration']
            proposals = np.asarray(proposals)

            # 64 * 6
            for i in range(proposals.shape[0]):
                current = np.copy(proposals[i,...])
                # current shape 6x4
                current[:,2] = 0
                pairs = mulity_scale_pairs(current)
                ret = []
                index = []
                for pos,pair in enumerate(pairs):
                    x = extract_DPN107_feature(feat,pair,duration,PYRAMID=[12,12],repeat=True)
                    if x is not None and x.shape[0] == 24:
                        ret.append(x)
                        index.append(pos)
                # TODO 兼容有问题
                # 由于特征提取的原因, 有时候并不能正好得到6张图片
                ret = np.asarray(ret)
                current = current[index,...]
                yield [vid,current,ret]

    def nextbatch(self):
        return self._random_sample_traing_item()

    def enum_genfucntion(self):
        return self._enum_every_test_item()

if __name__=='__main__':
    model = Dataset_DPN107_Pair('validation','train')
    x,y = model.nextbatch()
